//	Aviv Charikar	Yael Tzirulnikov	305343600	307955609

package modelTraining;

import java.io.*;
import java.util.*;

public class BasicModelTraining {
	private HashMap<String, HashMap<String, Integer>> bigram;
	private int bigramCount;
	private HashMap<String, Integer> unigram;
	private int unigramCount;

	public BasicModelTraining(String devFilename) throws IOException {
		bigram = new HashMap<String, HashMap<String, Integer>>();
		bigramCount = 0;
		unigram = new HashMap<String, Integer>();
		unigramCount = 0;
		readDevFile(devFilename);
	}
	
	private void readDevFile(String devFilename) throws IOException {
		BufferedReader reader = new BufferedReader(new FileReader(devFilename));
		String line;
		while ((line = reader.readLine()) != null) {
			// Handle only sentences
			if ((line.length() == 0) || (line.startsWith("<TRAIN")))
				continue;
			
			// Split the words in the sentence and add them to the counts
			String[] lineWords = line.split(" ");
			String lastWord = null;
			for (String word : lineWords) {
				if (word.length() > 0) {
					// Add the unigram
					Integer count = unigram.get(word);
					if (count == null)
						unigram.put(word, 1);
					else
						unigram.put(word, count + 1);
					unigramCount++;
					
					// Add the bigram
					if (lastWord != null) {
						HashMap<String, Integer> map = bigram.get(lastWord);
						if (map == null) {
							map = new HashMap<String, Integer>();
							bigram.put(lastWord, map);
						}
						count = map.get(word);
						if (count == null)
							map.put(word, 1);
						else
							map.put(word, count + 1);
						bigramCount++;
					}
					
					lastWord = word;
				}
			}
		}
		reader.close();
	}
	
	public HashMap<String, HashMap<String, Integer>> getBigrams() {
		return bigram;
	}
	
	public HashMap<String, Integer> getUnigrams() {
		return unigram;
	}	
	
	public int getBigramCount() {
		return bigramCount;
	}
	
	public int getUnigramCount() {
		return unigramCount;
	}
}
